"""
pytorch通用训练模板
==================
功能：适用于简单的回归和分类任务
作者：韩超明
创建时间：2025-10-23
最后更新：2025-10-23

使用说明：
1. 修改【数据准备】部分 -加载你自己的数据
2. 修改【模型定义】部分 - 设计你的网络结构
3. 调整【超参数配置】 - 根据任务调整参数
4. 运行训练
"""

import torch 
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm  # 进度条库
import time

# ========================================
# 第一部分：超参数配置
# ========================================

class Config:
    """配置类 - 把所有超参数集中管理"""

    #数据参数
    BATCH_SIZE = 32     # 批大小：每次训练使用多少个样本
    TRAIN_SPLIT = 0.8   # 训练集比例：80%用于训练，20%用于验证
    
    # 训练参数
    EPOCHS = 50              # 训练轮数：整个数据集训练多少遍
    LEARNING_RATE = 0.001    # 学习率：控制参数更新的步长
    
    #模型参数
    INPUT_SIZE = 2          #输入特征维度
    HIDDEN_SIZE = 16        #隐藏层神经元的数量
    OUTPUT_SIZE = 1         #输出维度

    #设备配置
    DEVICE = 'cpu'

    #其他
    RANDOM_SEED = 42      #随机种子：保证结果可重复
    SAVE_MODEL = True     #是否保存模型
    MODEL_PATH = 'best_model.pth'  #模型保存路径

# 设置随机种子，确保结果可重复
torch.manual_seed(Config.RANDOM_SEED)
np.random.seed(Config.RANDOM_SEED)

print(f"使用设备: {Config.DEVICE}")
print(f"PyTorch版本: {torch.__version__}")
print("-" * 50)


# ========================================
# 第二部分：数据集类定义
# ========================================

class CustomDataset(Dataset):
    """
    自定义数据集类
    
    继承自torch.utils.data.Dataset
    必须实现三个方法：__init__, __len__, __getitem__
    """
    def __init__(self,features,labels,transform=None):
        """
        初始化数据集
        
        参数：
            features: 特征数据 (numpy数组或tensor)
            labels: 标签数据 (numpy数组或tensor)
            transform: 数据转换函数 (可选)
        """       
        # 转换为tensor并移动到指定设备
        if isinstance(features, np.ndarray):
            self.features = torch.FloatTensor(features)
        else:
            self.features = features.float()
        
        if isinstance(labels, np.ndarray):
            self.labels = torch.FloatTensor(labels)
        else:
            self.labels = labels.float()
        
        self.transform = transform

        # 移动到设备
        self.features = self.features.to(Config.DEVICE)
        self.labels = self.labels.to(Config.DEVICE)

        print(f"数据集大小: {len(self.features)}")
        print(f"特征维度: {self.features.shape}")
        print(f"标签维度: {self.labels.shape}")

    def __len__(self):
        """
            返回数据集大小
            
            这个方法会被len()函数调用
        """ 
        return len(self.features)

    def __getitem__(self, idx):
        """
        根据索引返回单个样本
        
        参数：
            idx: 样本索引
        
        返回：
            (feature, label): 特征和标签的元组
        """
        feature = self.features[idx]
        label = self.labels[idx]

        # 如果有数据转换，应用转换
        if self.transform:
            feature = self.transform(feature)

        return feature, label


# ========================================
# 第三部分：模型定义
# ========================================
class NeuralNetwork(nn.Module):
    """
    神经网络模型
    
    网络结构：
    输入层 → 隐藏层1 → ReLU → Dropout → 隐藏层2 → ReLU → 输出层
    """

    def __init__(self, input_size, hidden_size, output_size, dropout=0.2):
        """
        初始化网络层
        
        参数：
            input_size: 输入特征维度
            hidden_size: 隐藏层神经元数量
            output_size: 输出维度
            dropout: Dropout比率（防止过拟合）
        """
        super().__init__()

        # 定义网络层
        self.layer1 = nn.Linear(input_size, hidden_size)  # 第一层：输入→隐藏
        self.layer2 = nn.Linear(hidden_size, hidden_size) # 第二层：隐藏→隐藏
        self.layer3 = nn.Linear(hidden_size, output_size) # 第三层：隐藏→输出

        # 激活函数
        self.relu = nn.LeakyReLUReLU()
        
        # Dropout层：训练时随机丢弃部分神经元，防止过拟合
        self.dropout = nn.Dropout(dropout)

        # 初始化权重（可选，使用Xavier初始化）
        self._initialize_weights()
    
    def _initialize_weights(self):
        """权重初始化"""
        for m in self.modules():
            if isinstance(m, nn.Linear):# isinstance()函数用于判断一个对象是否是一个已知的类型
                nn.init.xavier_uniform_(m.weight)# Xavier初始化：根据输入和输出维度，随机初始化权重，保持信号在前向传播和反向传播中的方差一致
                if m.bias is not None: 
                    nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        """
        前向传播
        
        参数：
            x: 输入数据
        
        返回：
            输出结果
        """
        # 第一层 + 激活 + Dropout
        x = self.layer1(x)
        x = self.relu(x)
        x = self.dropout(x)
        
        # 第二层 + 激活 + Dropout
        x = self.layer2(x)
        x = self.relu(x)
        x = self.dropout(x)
        
        # 输出层（回归任务不需要激活函数）
        x = self.layer3(x)
        
        return x

# ========================================
# 第四部分：训练和验证函数
# ========================================
def train_one_epoch(model, dataloader, criterion, optimizer, epoch):
    """
    训练一个epoch
    
    参数：
        model: 神经网络模型
        dataloader: 数据加载器
        criterion: 损失函数
        optimizer: 优化器
        epoch: 当前epoch数
    
    返回：
        平均损失值
    """
    model.train()  # 设置为训练模式（启用Dropout等）
    
    total_loss = 0.0
    num_batches = len(dataloader)
    
    # 使用tqdm显示进度条
    with tqdm(dataloader, desc=f'Epoch {epoch+1} [训练]') as pbar:
        for batch_idx, (data, target) in enumerate(pbar):
            # data和target已经在GPU上了（如果使用GPU）
            
            # 1. 前向传播
            output = model(data)
            loss = criterion(output, target)
            
            # 2. 反向传播
            optimizer.zero_grad()  # 清零梯度
            loss.backward()        # 计算梯度
            optimizer.step()       # 更新参数
            
            # 3. 记录损失
            total_loss += loss.item()
            
            # 4. 更新进度条
            pbar.set_postfix({'loss': f'{loss.item():.4f}'})
    
    # 返回平均损失
    avg_loss = total_loss / num_batches
    return avg_loss


def validate(model, dataloader, criterion, epoch):
    """
    验证函数
    
    参数：
        model: 神经网络模型
        dataloader: 数据加载器
        criterion: 损失函数
        epoch: 当前epoch数
    
    返回：
        平均损失值
    """
    model.eval()  # 设置为评估模式（禁用Dropout等）
    
    total_loss = 0.0
    num_batches = len(dataloader)
    
    # 不计算梯度，节省内存和计算
    with torch.no_grad():
        with tqdm(dataloader, desc=f'Epoch {epoch+1} [验证]') as pbar:
            for data, target in pbar:
                # 前向传播
                output = model(data)
                loss = criterion(output, target)
                
                # 记录损失
                total_loss += loss.item()
                
                # 更新进度条
                pbar.set_postfix({'loss': f'{loss.item():.4f}'})
    
    # 返回平均损失
    avg_loss = total_loss / num_batches
    return avg_loss


# ========================================
# 第五部分：主训练流程
# ========================================
def train_model(model, train_loader, val_loader, criterion, optimizer, 
                scheduler=None, num_epochs=Config.EPOCHS):
    """
    完整的训练流程
    
    参数：
        model: 神经网络模型
        train_loader: 训练数据加载器
        val_loader: 验证数据加载器
        criterion: 损失函数
        optimizer: 优化器
        scheduler: 学习率调度器（可选）
        num_epochs: 训练轮数
    
    返回：
        训练历史（包含每个epoch的损失）
    """
    # 记录训练历史
    history = {
        'train_loss': [],
        'val_loss': [],
        'learning_rates': []
    }
    
    # 记录最佳模型
    best_val_loss = float('inf')
    best_epoch = 0
    
    print("\n开始训练...")
    print("=" * 50)
    
    start_time = time.time()
    
    for epoch in range(num_epochs):
        # 记录当前学习率
        current_lr = optimizer.param_groups[0]['lr']
        history['learning_rates'].append(current_lr)
        
        # 训练
        train_loss = train_one_epoch(model, train_loader, criterion, 
                                     optimizer, epoch)
        history['train_loss'].append(train_loss)
        
        # 验证
        val_loss = validate(model, val_loader, criterion, epoch)
        history['val_loss'].append(val_loss)
        
        # 学习率调整
        if scheduler is not None:
            scheduler.step(val_loss)
        
        # 打印结果
        print(f"\nEpoch [{epoch+1}/{num_epochs}]")
        print(f"  训练损失: {train_loss:.6f}")
        print(f"  验证损失: {val_loss:.6f}")
        print(f"  学习率: {current_lr:.6f}")
        
        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_epoch = epoch + 1
            
            if Config.SAVE_MODEL:
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'loss': val_loss,
                }, Config.MODEL_PATH)
                print(f"  ✓ 保存最佳模型 (验证损失: {val_loss:.6f})")
        
        print("-" * 50)
    
    # 训练结束
    elapsed_time = time.time() - start_time
    print(f"\n训练完成！")
    print(f"总用时: {elapsed_time:.2f}秒")
    print(f"最佳模型: Epoch {best_epoch}, 验证损失: {best_val_loss:.6f}")
    
    return history


# ========================================
# 第六部分：结果可视化
# ========================================
def plot_training_history(history):
    """
    绘制训练历史曲线
    
    参数：
        history: 训练历史字典
    """
    fig, axes = plt.subplots(1, 2, figsize=(15, 5))# 1行2列的子图
    
    # 绘制损失曲线
    axes[0].plot(history['train_loss'], label='训练损失', linewidth=2)
    axes[0].plot(history['val_loss'], label='验证损失', linewidth=2)
    axes[0].set_xlabel('Epoch', fontsize=12)
    axes[0].set_ylabel('Loss', fontsize=12)
    axes[0].set_title('训练和验证损失', fontsize=14, fontweight='bold')
    axes[0].legend(fontsize=10)
    axes[0].grid(True, alpha=0.3)
    
    # 绘制学习率变化
    axes[1].plot(history['learning_rates'], color='orange', linewidth=2)
    axes[1].set_xlabel('Epoch', fontsize=12)
    axes[1].set_ylabel('Learning Rate', fontsize=12)
    axes[1].set_title('学习率变化', fontsize=14, fontweight='bold')
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('training_history.png', dpi=300, bbox_inches='tight')
    print("\n训练曲线已保存到 'training_history.png'")
    plt.show()


# ========================================
# 第七部分：模型测试和预测
# ========================================
def test_model(model, test_data, test_labels):
    """
    测试模型性能
    
    参数：
        model: 训练好的模型
        test_data: 测试数据
        test_labels: 测试标签
    """
    model.eval()
    
    with torch.no_grad():
        # 转换为tensor
        if isinstance(test_data, np.ndarray):
            test_data = torch.FloatTensor(test_data).to(Config.DEVICE)
        if isinstance(test_labels, np.ndarray):
            test_labels = torch.FloatTensor(test_labels).to(Config.DEVICE)
        
        # 预测
        predictions = model(test_data)
        
        # 计算误差
        mse = nn.MSELoss()(predictions, test_labels)
        mae = torch.mean(torch.abs(predictions - test_labels))
        
        print("\n" + "=" * 50)
        print("测试结果：")
        print(f"  均方误差 (MSE): {mse.item():.6f}")
        print(f"  平均绝对误差 (MAE): {mae.item():.6f}")
        print("=" * 50)
        
        # 显示几个预测示例
        print("\n预测示例（前5个）：")
        print("-" * 50)
        for i in range(min(5, len(test_data))):
            print(f"样本 {i+1}:")
            print(f"  输入: {test_data[i].cpu().numpy()}")
            print(f"  真实值: {test_labels[i].item():.4f}")
            print(f"  预测值: {predictions[i].item():.4f}")
            print(f"  误差: {abs(predictions[i].item() - test_labels[i].item()):.4f}")
            print()
    
    return predictions


# ========================================
# 第八部分：数据准备示例
# ========================================
def prepare_sample_data():
    """
    准备示例数据
    
    这里使用简单的线性关系：y = 2*x1 + 3*x2 + 噪声
    
    返回：
        X_train, y_train, X_val, y_val, X_test, y_test
    """
    print("\n生成示例数据...")
    
    # 生成数据
    n_samples = 1000
    n_features = 2
    
    # 特征：在[-10, 10]范围内随机生成
    X = np.random.uniform(-10, 10, (n_samples, n_features))
    
    # 标签：y = 2*x1 + 3*x2 + 噪声
    y = 2 * X[:, 0] + 3 * X[:, 1] + np.random.normal(0, 0.5, n_samples)
    y = y.reshape(-1, 1)
    
    # 划分数据集
    n_train = int(n_samples * Config.TRAIN_SPLIT)
    n_val = int(n_samples * 0.1)
    n_test = n_samples - n_train - n_val
    
    X_train, y_train = X[:n_train], y[:n_train]
    X_val, y_val = X[n_train:n_train+n_val], y[n_train:n_train+n_val]
    X_test, y_test = X[n_train+n_val:], y[n_train+n_val:]
    
    print(f"  训练集: {len(X_train)} 样本")
    print(f"  验证集: {len(X_val)} 样本")
    print(f"  测试集: {len(X_test)} 样本")
    
    return X_train, y_train, X_val, y_val, X_test, y_test


# ========================================
# 第九部分：主程序
# ========================================
def main():
    """主程序入口"""
    
    print("\n" + "=" * 50)
    print("PyTorch训练模板 - 开始运行")
    print("=" * 50)
    
    # 1. 准备数据
    X_train, y_train, X_val, y_val, X_test, y_test = prepare_sample_data()
    
    # 2. 创建数据集和数据加载器
    train_dataset = CustomDataset(X_train, y_train)
    val_dataset = CustomDataset(X_val, y_val)
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=Config.BATCH_SIZE,
        shuffle=True,  # 训练时打乱数据
        num_workers=0  # Windows上设为0，Linux/Mac可以设为2-4
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=Config.BATCH_SIZE,
        shuffle=False,  # 验证时不打乱
        num_workers=0
    )
    
    print(f"\n数据加载器创建完成")
    print(f"  训练批次数: {len(train_loader)}")
    print(f"  验证批次数: {len(val_loader)}")
    
    # 3. 创建模型
    model = NeuralNetwork(
        input_size=Config.INPUT_SIZE,
        hidden_size=Config.HIDDEN_SIZE,
        output_size=Config.OUTPUT_SIZE
    ).to(Config.DEVICE)
    
    print(f"\n模型结构:")
    print(model)
    
    # 统计参数数量
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\n总参数量: {total_params:,}")
    print(f"可训练参数: {trainable_params:,}")
    
    # 4. 定义损失函数和优化器
    criterion = nn.MSELoss()  # 回归任务用MSE，分类任务用CrossEntropyLoss
    
    optimizer = optim.Adam(
        model.parameters(),
        lr=Config.LEARNING_RATE,
        weight_decay=1e-5  # L2正则化，防止过拟合
    )
    
    # 5. 学习率调度器（可选）
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='min',      # 监控指标降低时调整
        factor=0.5,      # 学习率衰减因子
        patience=5,      # 等待epochs数
        verbose=True     # 打印信息
    )
    
    # 6. 开始训练
    history = train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        criterion=criterion,
        optimizer=optimizer,
        scheduler=scheduler,
        num_epochs=Config.EPOCHS
    )
    
    # 7. 可视化训练过程
    plot_training_history(history)
    
    # 8. 加载最佳模型并测试
    if Config.SAVE_MODEL:
        print("\n加载最佳模型...")
        checkpoint = torch.load(Config.MODEL_PATH)
        model.load_state_dict(checkpoint['model_state_dict'])
        print("✓ 最佳模型加载完成")
    
    # 9. 测试模型
    test_model(model, X_test, y_test)
    
    # 10. 保存完整模型（可选）
    torch.save(model, 'complete_model.pth')
    print("\n完整模型已保存到 'complete_model.pth'")
    
    print("\n" + "=" * 50)
    print("程序执行完毕！")
    print("=" * 50)


# ========================================
# 运行程序
# ========================================
if __name__ == "__main__":
    main()
